open Utils.RandSelect
open Utils.Pervasive
open Utils.ListSpecial
open Utils.StringSpecial

open Syntax
open Binding
open PrintExpr
open GPUtils

open BlockPopulation.Gene 
open BlockPopulation

open BlockPopulationVisual
open Expr_visual
open Pointers
open TestCases
open ProofAssembler
open Population
open ProofVisual





let kill_g population genotype = 
		match Hashtbl.find population.genotypes genotype with
		|Alive ev -> Hashtbl.replace population.genotypes genotype (Dead ev)
		|_->()
let kill eco species genotype = 
		kill_g (Hashtbl.find eco.kingdom species) genotype
 
let iter_population ?(pred=(fun _-> true)) ?(f=(fun _ _ -> ( )))   population = 
	Hashtbl.iter
	(fun genotype status -> if pred status then f genotype status)
	population	
	
let selection_filter_genotypes eco threshold = 
	Hashtbl.iter 
	(fun proof pop -> iter_population  pop.genotypes ~pred:(function Alive _ ->true |_ ->false) 
	~f:(fun prog status ->
	match status with  
	| Alive  x when x <= threshold -> Hashtbl.replace pop.genotypes prog (Dead x)
	|_ ->( ))) eco.kingdom
	
let rec reap_genotypes eco = 
	match card_ecosystem ~pred:(function Alive _ -> true | _->false) eco with
		| x when x > eco.settings.aprNumberOfLiveGenotypes - eco.settings.selectReap  ->
			let species = List.hd (choose_random ~selectionSize:1 (keys eco.kingdom)) in
			let live_population = get_live (Pop (Hashtbl.find eco.kingdom species)) in
			(match List.length live_population with
					| x when x > 0 ->
						prf ("\nReaping:" ^ (string_of_int x));
						let (ev, genotype) = 
							List.hd (rouletteWheelSelect live_population (fun (a,_) -> if a > 0. then 1. /. a else 0.)) in 
						kill eco species genotype;
						reap_genotypes eco
					| _-> reap_genotypes eco)				
							
		| _ -> ()	

let selection_filter_species eco =
	Hashtbl.iter 
	(fun species population -> 
		match card_population ~pred:(function Alive _->true | _->false) population with
		| x when x <= 0 -> Hashtbl.remove eco.kingdom species
		| card_live_population -> 
			let sum = ref 0. in
			Hashtbl.iter (fun _ -> function Alive x -> sum := !sum +. x |_->()) population.genotypes;
			population.info.avgFitness <- (!sum /.  (float card_live_population))) 
	eco.kingdom 
		
let selection_filter_genes eco	 =
	let genes = !(eco.genes) in
	let count_gene_use eco =
		reset_values genes;		
		let rec update species prog ev = 
			match (species, prog) with 
			| ProgNode ty, In id -> 
				let catInfo, lst = tyFind genes ty in
				catInfo.species_use <- catInfo.species_use + 1;
				let alle = List.find (alleHasId id) lst in
					alleIncUseInPop alle;
					alleSetFitness alle ((alleFitness alle) +. ev)
			| ForAllElim (proof, ty1), TypeFunc (prog, ty2) -> 
				update proof prog ev
			| (ImplyElim (proof1, proof2)) , (ObjFunc (prog1, prog2)) ->
				update proof1 prog1 ev; update proof2 prog2 ev
			| _->failwith "Invalid pair proof/program" in
		Hashtbl.iter (fun species population ->
			Hashtbl.iter (fun prog -> function 
					| Alive ev -> update species prog ev;
					|_->( )) population.genotypes) eco.kingdom in
	 
	 count_gene_use eco;
	 gpRemoveUnusedAlleles genes
	
		
let rec update_population_info eco species population =
	let genes = !(eco.genes) in
	let population_info = population.info in
		population_info.fillProportion <- 0.;
		population_info.theoreticalMaxFill <- 0;
		population_info.currentFill <- 0;
		population_info.avgComplexity <- 0.;
		population_info.avgFitness <- 0.;		
	let genotypes = population.genotypes in
	let genos = progs_of_proof genes species in 
	let geno_status_lst = List.map (fun genotype -> 
		match (Hashtbl.mem genotypes genotype) with
		| false -> (genotype, Unborn)
		| true -> (genotype, Hashtbl.find genotypes genotype)) genos in
	
	population_info.theoreticalMaxFill <- List.length geno_status_lst;
	let new_gene_status_hash = Hashtbl.create population_info.theoreticalMaxFill in
	List.iter (fun (geno,status) -> 
		(match status with 
		|Alive fitVal -> 
			population_info.currentFill <- population_info.currentFill + 1;
			population_info.avgFitness <- population_info.avgFitness +. fitVal
		|_->());
		Hashtbl.add new_gene_status_hash geno status) geno_status_lst;
	if population_info.currentFill > 0 then(	 
		population_info.avgFitness <- population_info.avgFitness /.	(float) population_info.currentFill;
		population_info.fillProportion <- (float) population_info.theoreticalMaxFill /. (float) population_info.currentFill);
	Hashtbl.replace (eco.kingdom) species 
		{
			info = population_info;
			contains_solution = false;
			genotypes = new_gene_status_hash
		} 	
let liveGenotypes eco = card_ecosystem ~pred:(function Alive _ -> true| _-> false) eco	

let next_generation ?(debug=false) eco constraintFunc genotypeEvaluationNum = 
	
	let genes = !(eco.genes) in
	let ctx = gpContext genes  in
	let non_determinisic_selection_threshold = 0. in
	selection_filter_genotypes eco non_determinisic_selection_threshold;
	reap_genotypes eco;
	
	let clean_ecosystem( ) = 
		selection_filter_genotypes eco non_determinisic_selection_threshold;
		selection_filter_species eco;
		selection_filter_genes eco;
		Hashtbl.iter (update_population_info eco)  eco.kingdom in 
			
	let ecosystemIsNotFull( ) = 
		(liveGenotypes eco < eco.settings.aprNumberOfLiveGenotypes) in
	
		
			
	match liveGenotypes eco with
		| 0 -> (*No genotypes survived, start over*)
			None(*init_generation eco constraintFunc genotypeEvaluationNum*)
		| card_alive_ecosystem ->
			clean_ecosystem ( );
			let foundSolution = ref false in
			let sol = ref None in
			(*while (ecosystemIsNotFull( ) && not(!foundSolution)) do
				match addNewLiveGenotypes( ) with
				|None -> ()
				|Some solution -> 
					foundSolution := true;
					sol := Some solution
			done;*)		 		
			eco.generationCounter <- eco.generationCounter + 1;
			!sol   
		
			
		











open List
open Gobject.Data
open Core
open StdLabels
open GtkTree


let cols = new GTree.column_list
let max_inds = cols#add string
let lives = cols#add string
let deads = cols#add string
let unborns = cols#add string
let aidCol = cols#add string
let specFitCol = cols#add string
let expr_species = cols#add string
let bg = cols#add (unsafe_boxed (Gobject.Type.from_name "GdkColor"))

let create_model ( ) =
  let model = GTree.tree_store cols in
  model
let add_columns ~(view : GTree.view) ~model =
  let renderer = GTree.cell_renderer_text [`XALIGN 0.] in
  List.iter 
  [("Avg Fitness", specFitCol);
  ("Max" , max_inds) ;  
  ("Alive" , lives) ;
  ("Dead", deads) ;
  ("Unborn", unborns); 
  ("Species ID" , expr_species);
  ("Species Ancestor ID", aidCol)]
  
  ~f: (fun (title, col) -> 
			let col_view = GTree.view_column ~title
			~renderer:(renderer, [("text", col)]) ( ) in
			col_view#add_attribute renderer "background-gdk" bg;
			un(view#append_column col_view))			
type modelSettings = ShowIncompleteProofs 
| SortProofs of (speciesInfo -> speciesInfo -> int) 
| SolutionHandle of (ecosystem->unit)
let updateModel ?(debug=false) eco ~(model : GTree.tree_store) settings = 
		let g_pop = !(eco.genes) in
		let ctx = gpContext g_pop in
		begin try let iter = model#iter_children ~nth:0 None in
			let remove = ref (model#remove iter) in
			while !remove do
				remove := model#remove iter 
			done
		with _->( ) end;
		
		
		if List.exists (function SolutionHandle _ -> true |_->false) settings
		then(
			if List.length eco.solutions > 0 then (
			match (find (function SolutionHandle _ -> true |_->false) settings)
			with SolutionHandle handle -> handle eco
			|_->raise Not_found)
			 (*gene_window eco.genes 
			 [	ShowOnlySubset (List.map (fun te -> Te(te, None)) eco.solutions);
				DisplayEvaluate ])*));		
		let spec_pop_pairs = Hashtbl.fold 
		(fun species pop acc -> (species,pop) :: acc) eco.kingdom [] in
		
		let spec_pop_pairs = 
			if List.exists (function ShowIncompleteProofs->true |_->false) settings 
			then spec_pop_pairs
			else List.filter (fun (species,_) -> is_complete species) spec_pop_pairs in
			
		let spec_pop_pairs = 
			if List.exists (function SortProofs _ ->true |_->false) settings 
			then (
				let sort = 
					match find (function SortProofs _ ->true |_->false) settings
					with SortProofs pf -> (fun (_, pop1) (_, pop2) -> 
						pf pop1.info pop2.info)
					| _-> raise Not_found in
				List.fast_sort sort spec_pop_pairs)
			else spec_pop_pairs in
		
		ListLabels.iter spec_pop_pairs
		~f:(fun (species , population) ->
			let row = model#append ( ) in 
			let set column = model#set ~row ~column in 
			let color str = set bg (GDraw.color (`NAME str)) in
			
			if population.contains_solution 
				then color "yellow"
				else color "blue";
			
			List.iter ~f:(fun (str, column) -> model#set ~row ~column str)
			([(string_of_float (population.info.avgFitness), specFitCol);
			((string_of_int (population.info.speciesId)) , expr_species)  ;
			((match population.info.ancesterID with 
				|None -> "--"
				| Some id ->string_of_int id), aidCol)] @
		
			(List.map (fun (pred, col) -> string_of_int (card_population ~pred population), col) 
			[((fun _ -> true), max_inds); 
			((function Alive _ -> true |_->false), lives);
			((function Dead _ -> true |_->false), deads);
			((function Unborn  -> true |_->false), unborns)]));					
			Hashtbl.iter (fun genotype genStat ->
				match genStat with 
				| Alive _ ->
				(let row = model#append ~parent:row () in 
				let set column = model#set ~row ~column in 
				let color str = set bg (GDraw.color (`NAME str)) in
				(match genStat with 
					| Alive ev when ev = eco.maxGenotypeScore ->
						color "red";
						set lives (string_of_float ev)
					| Alive ev -> 
						color "orange";
						set lives (string_of_float ev)
					| _->( ));
					set expr_species (string_of_genotype ~separator:":" g_pop genotype))
				|_->( )) population.genotypes) 
			
			
			

(*let prog_of_string ?(debug=false) str = 
	let reg =  Str.regexp "[- ]*\\(([0-9]+)\\)\\(.*\\)" in
	if Str.string_match reg str 0 then (
		let str1 = Str.matched_group 1 str in
		let str2 = Str.matched_group 2 str in 
		let prog = [reference_of_string str1] in
		let rec fill prog str0 = 
			if Str.string_match reg str0 0 then 
				(let str1 = Str.matched_group 1 str0 in
				 let str2 = Str.matched_group 2 str0 in
				 fill (prog @ [reference_of_string str1]) str2)
			else prog in fill prog str2 )
		else failwith "Unable to produce program from string" *)


let on_row_activated ?(debug=false) eco (view:GTree.view) path column = 
	let model = view#model in
	let row = model#get_iter path in
	let str = model#get ~row ~column:expr_species in
	try 
		let id = int_of_string str in
		let species = get_species_from_id !eco id in
		prf ~filename:(Some "tree.dot") (string_of_proof  ~output:"dot" (gpContext !(!eco.genes)) species);
		command "dot -Tgif tree.dot -o out.gif | out.gif"
	with
	|_->( )(*
	let g_pop = !(!gp_ref.genes) in
	let ctx = !(g_pop.ctx) in
	 in
	try
		let prog = prog_of_string ~debug str in
		let te = term_of_program ~debug g_pop prog in 
			show_expr_window ctx (Te(te, None)) ( )
	with _-> prf str	*)
		
		
(*Test cases  and GoalType loading ****************************)			
let openLast = ref None
let testCases_filename = ref None

let leftCombinatorLst ctx = 
	List.map
	(fun str -> ter (parse_string str ctx))
	["lam a: TyBool->TyBool->TyBool . lam o: TyBool->TyBool->TyBool . lam x : TyBool . lam y : TyBool . Lam Y . lam z : TyBool->TyBool->Y. z (a x y) (o x y)"]

let readCases ps_file ~(gp : ecosystem ref) ~(gpSettings:ecoSetting) callb =
	let ctxRef = ref (gpContext !(!gp.genes)) in
	let (cases, ty) = read_test_cases ps_file !ctxRef in
	gp := make_eco (ref (BlockPopulation.make ctxRef ty (leftCombinatorLst !ctxRef) [])) ty cases gpSettings;
	(match !openLast with
	|Some fsett -> 
		let oc = open_out fsett in
		output_string oc (Filename.dirname ps_file ^ "\\\n"); close_out oc
	|_->( ));
	callb ps_file 
let file_ok_sel filew ~(gp : ecosystem ref) ~(gpSettings:ecoSetting) callb ( )= 	
	let ps_file =  filew#filename in
	readCases ps_file ~gp ~gpSettings callb;
	filew#destroy( )
	
	
let load_tests ?(filename = "") ~(gp:ecosystem ref) ~(gpSettings:ecoSetting) callb ( ) =
	let filew = if filename= "" 
		then GWindow.file_selection 
			~title:"Select test case file" 
			~border_width:10  ( ) 
		else GWindow.file_selection 
			~title:"Select test case file" 
			~border_width:10 ~filename ( ) in
	
	un (filew#ok_button#connect#clicked 
		~callback:(file_ok_sel filew ~gp ~gpSettings callb));
	un (filew#cancel_button#connect#clicked ~callback:filew#destroy);
	filew#show()    



let empty_goal_type ctx =	
	parse_string "(TT X. X -> X)->Int" ctx 

		
type population_window_setting = Auto | ProofDepth of int | BlocPopInitComplexity of int 
| LoadCasesOnOpen of string | OpenFileDlgPathGP of string | GenePrecision of int 
|PopSize of int | MaxGeneSize of int | MaxProofWidth of int			
let make_population_visual ?(debug=false) refCtx settings =
	match !refCtx with 
	| [] -> failwith "Can not start GP until initial context is loaded\n" 
	| _  -> 
	
	(*Window settings***************)
	let window = GWindow.window ~title:"population" () in
		window#set_default_size ~width:950 ~height:700;
    
	let hBox = GPack.hbox ~packing:window#add () in
	let leftPanel = GPack.vbox 
		~border_width:8 
		~height:350 
		~width:500 
		~spacing:8 
		~packing:hBox#add () in    
    
    (*space for program listing *)
	let sw_progs, model = 
	GBin.scrolled_window 
		~shadow_type:`ETCHED_IN ~hpolicy:`AUTOMATIC 
		~vpolicy:`AUTOMATIC ~packing:leftPanel#add (), 
	create_model ( ) in
    
    let view = GTree.view ~model:model ~packing:sw_progs#add () in 
	view#set_rules_hint true; 
	view#selection#set_mode `MULTIPLE; 
	add_columns ~view ~model;
    let boxMenu = GPack.vbox ~width:80 ~packing:hBox#add () in
	let panel ?(box = boxMenu) title = 
		let pan = GBin.frame ~label:(title^ "...") ~packing:box#pack ( ) in
		GPack.vbox ~width:80 ~packing:pan#add ( ) in
	
	let controlPanel = panel "Tools..." in
	let genePoolPanel = panel "Gene pool..." in
	let testCasePanel =  panel "Test cases..." in     
    let proofPanel =  panel "Proofs" in 
	let gpPanel =  panel "gp..." in 
    	
    let empty_entry ?(editable=false) ?(text="") packer ( ) = 
    GEdit.entry ~editable ~text  ~width:100  ~packing:packer#add ( ) in
     
    let button str packer ( ) = GButton.button ~label:str ~packing:packer#add () in
	let update_btn = button "Update labels" controlPanel ( ) in
	
    let num_elt_label, tot_complexity_label, 
    precision_label, max_size_label, open_blocks_btn  = 
	empty_entry genePoolPanel ( ),
	empty_entry genePoolPanel ( ),
    empty_entry genePoolPanel ( ),
    empty_entry genePoolPanel ( ),
    button "View blocks" genePoolPanel ( ) in 

    let load_cases_button, view_cases_button, 
    file_name_label, goal_type_label = 
	button "Load test cases" testCasePanel ( ), 
	button "View/Edit cases" testCasePanel ( ), 
	empty_entry testCasePanel ( ),
	empty_entry testCasePanel ( ) in 
    
        
    let proof_num_label, proofView_button, prog_label = 
    empty_entry ~text:"No proofs" proofPanel ( ),
    button "View selected proof" proofPanel ( ),
	empty_entry ~text:"No programs possible" proofPanel ( ) in
    let show_all_proofs = GButton.check_button  
				~label:"Show all" 
				~active:false
				~packing:proofPanel#add ( ) in  

	(*Parameters  ******************************)
	let genePoolStepComplexity, precision, 
	proofDepth, popSize, mgSize, 
	maxProofWidth, auto = 
	(fold_left (fun default_val -> function
		BlocPopInitComplexity comp->comp
		|_->default_val) 10000 settings), 
	(fold_left (fun default_val -> function
		GenePrecision pre  -> pre
		|_->default_val) 5 settings),
	(fold_left (fun default_val -> function
		ProofDepth pd -> pd
		|_->default_val) 3 settings),
	(fold_left (fun default_val -> function
		PopSize pd -> pd
		|_->default_val) 1000 settings),
	(fold_left (fun default_val -> function
		MaxGeneSize mg ->mg
		|_->default_val) 30 settings),
	(fold_left (fun default_val -> function
		MaxProofWidth mpw -> mpw
		|_->default_val) 5 settings),
	List.exists (function Auto -> true |_->false) settings in   
    let gpSettings = {
			maxProofWidth = maxProofWidth;
			aprGenePoolTotComp=genePoolStepComplexity;
			aprNumberOfLiveGenotypes=popSize;
			selectSize = 40;
			selectReap = 40 } in
	let constrFunc = (fun _ expr-> complexity_of_expr expr < mgSize) in
    let population_live_label, population_dead_label, generationLabel = 
	empty_entry  gpPanel ( ),
	empty_entry  gpPanel ( ), 
	empty_entry  gpPanel ( ) in  
	 
	let refGenes = ref (BlockPopulation.make refCtx (typ (empty_goal_type !refCtx)) [] []) in 
	let gp = ref (empty_eco refGenes) in
	let sw_stats = 
	GBin.scrolled_window 
		~shadow_type:`ETCHED_IN ~hpolicy:`AUTOMATIC 
		~vpolicy:`AUTOMATIC ~packing:gpPanel#add () in
	let stats_text = GText.view 
			~width:350 
			~height:200 
			~border_width:3 
			~wrap_mode:`WORD
			~packing:sw_stats#add( ) 	in 	
	let update ( ) = 
		let genes = !(!gp.genes) in
		let ctx = gpContext genes in
		
		List.iter (fun (a,b) -> a#set_text b)
		[(generationLabel, ("Generation " ^ string_of_int !gp.generationCounter));
		(population_live_label, ("Evaluated programs : " ^ 
		string_of_int (card_ecosystem ~pred:(function Alive _ ->true |_->false) !gp)))];
	
		goal_type_label#set_text ("Goal type is " ^ string_of_type ctx !gp.goalType Compact);
		let proofs = (keys !gp.kingdom) in
			(proof_num_label#set_text 
				("There are " ^ string_of_int (List.length proofs) ^ " proofs in the system"));
		prog_label#set_text 
				("There are " ^ 
				string_of_int (prog_num genes proofs) ^ " possible programs");
		
		stats_text#buffer#set_text !gp.stats;	
		updateModel ~debug !gp ~model 
		(if show_all_proofs#active then [ShowIncompleteProofs] else [ ]);
		
		
		List.iter2 (fun a b -> a#set_text b)
		[num_elt_label;tot_complexity_label;max_size_label]
		["There are " ^ string_of_int (card  !(!gp.genes)) ^ " elts in the pool";
		"Total complexity is " ^ string_of_int (comp !(!gp.genes)) ;
		(*"The precision is at " ^ string_of_int !refGene_pool.terms.precision ; *)
		"The max size of a spawned gene is " ^ string_of_int mgSize] in
	
	un(show_all_proofs#connect#toggled  ~callback:update);
	let callb filename = 
		testCases_filename := Some filename;
		file_name_label#set_text filename;
		update( ) in		
	
		
    (try match find (function  OpenFileDlgPathGP _->true |_->false) settings 
	with OpenFileDlgPathGP filename ->
			let filename = 
			(try (let ic = open_in (filename ^ "gpWindowSettings.dat") in 	 
				let res = input_line ic in close_in ic;
				openLast := Some (filename ^ "gpWindowSettings.dat"); res)
				with _ -> filename) in
			load_cases_button#connect#clicked 
				~callback:(load_tests ~filename ~gp ~gpSettings callb)
    		|_->raise Not_found
	with _-> load_cases_button#connect#clicked 
					~callback:(load_tests ~gp ~gpSettings callb));
    (try match find (function LoadCasesOnOpen _ ->true |_->false) settings with
		| LoadCasesOnOpen filename -> readCases filename ~gp ~gpSettings callb
		| _->raise Not_found
	with |_-> ( ));
	
    let view_cases ( ) = 
		match !testCases_filename with
		| None->( )
		| Some filename -> command filename in
	un(view_cases_button#connect#clicked ~callback:view_cases);	
    let show_proofs( ) = 
		let selection = view#selection in
		let selectedRows = selection#get_selected_rows in
		let path = List.hd selectedRows in
		let row = model#get_iter path in
		let str = model#get ~row ~column:expr_species in
		try 
			let id = int_of_string str in
			let species = get_species_from_id !gp id in
			make_proof_visual !gp (id, species)
		with
		|_->( ) in	
    un(proofView_button#connect#clicked ~callback:show_proofs);
    let show_genes ( ) =
		gene_window !gp.genes
		[ShowGrowByOne ; 
		ShowGrowByComp genePoolStepComplexity ; 
		MaxSize mgSize] in
	un(open_blocks_btn#connect#clicked ~callback:show_genes);
    
   
	 let init_first_generation ( ) = 
		let genCtr = ref 0 in
		let gen_func = match !gp.generationCounter with
		| 0 -> init_generation
		| _ -> next_generation in
		prf ("\nResults: " ^ 
		(match gen_func !gp constrFunc genCtr with 
		|Some _ -> "Found solution"
		|_-> "Didn\'t find solution"));
	prf ("\nTotal of " ^ (string_of_int !genCtr) ^ " evaluations");
	update ( ) in	
	(*Button connects ***********************************************)
	un(update_btn#connect#clicked ~callback:update);
	un(view#connect#row_activated ~callback:(on_row_activated ~debug gp view));
    
  
	
	if auto 
	then begin
		gene_window !gp.genes
		[Hide; BuildInit genePoolStepComplexity; 
			MaxSize mgSize];
		init_first_generation ( ) end
	else begin
	let init_first_generation_btn = 
	button "Next generation" gpPanel ( ) in 
	un(init_first_generation_btn#connect#clicked ~callback:(init_first_generation))
	end;
	window#show()
